#!/usr/bin/python
from rap_core import *


zeff=8
core=-110.
ip=0.5


clevs=rap_states_lsj()
clevs.set_nuclear_charge(zeff)
clevs.set_core_energy(core)
clevs.set_ionization_threshold(ip)

clevs.add_level('1s2','1S',0,0.0)
clevs.add_level('1s 2s','3S',2,41.231553)
clevs.add_level('1s 2p','3P',0,41.787244)
clevs.add_level('1s 2p','3P',2,41.787791)
clevs.add_level('1s 2p','3P',4,41.792803)
clevs.add_level('1s 2s','1S',0,41.812395)
clevs.add_level('1s 2p','1P',2,42.184383)

if clevs.nuclear_charge() != zeff:
  print 'problem with nuclear_charge()'
if clevs.core_energy() != core:
  print 'problem with core_energy()'
if clevs.ionization_threshold() != ip:
  print 'problem with ionization_threshold()'
if len(clevs) != 7:
  print 'wrong number of levels'

if str(clevs[1]) != '1s2 1S_0':
  print 'A - incorrect level: 1s2 1S_0'
if str(clevs[5]) != '1s1 2p1 3P_2':
  print 'A - incorrect level: 1s 2p 3P_2'


toadd=[]
toadd.append(( 0.000000,2,'1s2','1S',0))
toadd.append((41.231553,3,'1s 2s','3S',2))
toadd.append((41.787244,4,'1s 2p','3P',0))
toadd.append((41.787791,5,'1s 2p','3P',2))
toadd.append((41.792803,6,'1s 2p','3P',4))
toadd.append((41.812395,7,'1s 2s','1S',0))
toadd.append((42.184383,8,'1s 2p','1P',2))


clevs2=clevs.copy()
clevs2.clear_levels()
clevs2.add_levels(toadd)

if clevs2.nuclear_charge() != zeff:
  print 'problem with nuclear_charge()'
if clevs2.core_energy() != core:
  print 'problem with core_energy()'
if clevs2.ionization_threshold() != ip:
  print 'problem with ionization_threshold()'
if len(clevs2) != 7:
  print 'wrong number of levels'

if str(clevs2[2]) != '1s2 1S_0':
  print 'B - incorrect level: 1s2 1S_0'
if str(clevs2[6]) != '1s1 2p1 3P_2':
  print 'B - incorrect level: 1s 2p 3P_2'

clevs.set_lmap(clevs2)
if clevs.label(3,imap=True,spacer='.') != '1s1.2s1.3S_1':
  print 'problem with mapping or label'

clevs2.reindex()
if str(clevs2[1]) != '1s2 1S_0':
  print 'C - incorrect level: 1s2 1S_0'
if str(clevs2[5]) != '1s1 2p1 3P_2':
  print 'C - incorrect level: 1s 2p 3P_2'


levs=rap_states_lsj()
levs.load_olg('olg.v21')
idxc=levs.cfilter('2s 2p3')
idxt=levs.tfilter('3P',subset=idxc)
idx2j=levs.j2filter(4,subset=idxc)
idxj=levs.jfilter(2,subset=idxc)
if idxc != [5,7,8,9,10,11,12,13,14,15]:
  print 'problem with cfilter()'
if idxt != [10,11,12]:
  print 'problem with tfilter()'
if idx2j != [5,8,12,13]:
  print 'problem with tfilter()'
if idxj != [5,8,12,13]:
  print 'problem with tfilter()'


levs=rap_states_lsj('calike_mcw07#fe6.dat')
if len(levs) != 189:
  print 'problem opening adf04 file'
